# import argparse
# import threading
# import warnings
# import requests
# from mysql_val import mysql_run
# from redis_val import redis_run
# from ip_val import run_ip
# from ssh_val import ssh_run  # 导入SSH漏扫功能
# from ccolor import Colors
#
# warnings.filterwarnings("ignore")
#
# banner = """
#         __                        .__
#   _____|  | ___________  _______  |  |
#  /  ___/  |/ / ____/\  \/ /\__  \ |  |
#  \___ \|    < <_|  | \   /  / __ \|  |__
# /____  >__|_ \__   |  \_/  (____  /____/
#      \/     \/  |__|            \/
# """
# print(f"{Colors.BLUE_BRIGHT}{Colors.BOLD}{banner}{Colors.RESET}")
#
# headers = {
#     "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/89.0.4389.90 Safari/537.36",
# }
#
# def subdomain_scan(domain, sub):
#     host1 = "http://" + sub + "." + domain
#     host2 = "https://" + sub + "." + domain
#     try:
#         response1 = requests.get(host1, headers=headers, timeout=5)
#         if response1.status_code in [200,301,302,401,403,404]:
#             print(f"[+] {sub}.{domain} 存活")
#     except:
#         try:
#             response2 = requests.get(host2, headers=headers, timeout=5)
#             if response2.status_code in [200,301,302,401,403,404]:
#                 print(f"[+] {sub}.{domain} 存活")
#         except:
#             pass
#
# def run_domain(domain, sub_file="dic.txt"):
#     threads = []
#     with open(file=sub_file, mode="r", encoding="utf-8") as f:
#         for sub in f:
#             sub = sub.strip()
#             t = threading.Thread(target=subdomain_scan, args=(domain, sub))
#             t.start()
#             threads.append(t)
#     for t in threads:
#         t.join()
#
# # 创建命令行解析器对象
# parser = argparse.ArgumentParser(description="智慧AI漏扫")
# subparsers = parser.add_subparsers(help="选择要使用的功能", dest="sub_command")
#
# # 创建MySQL的解析器对象
# mysql_subparsers = subparsers.add_parser("mysql", help="MySQL相关漏扫")
# mysql_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="MySQL主机地址")
# mysql_subparsers.add_argument("--port", "-P", dest="port", type=int, default=3306, help="MySQL端口号")
# mysql_subparsers.add_argument("--username", "-u", dest="username", type=argparse.FileType(mode='r', encoding="utf-8"),
#                               default="./user.dict", help="MySQL用户名")
# mysql_subparsers.add_argument("--password", "-p", dest="password", type=argparse.FileType(mode='r', encoding="utf-8"),
#                               default="./passwd.dict", help="MySQL密码本")
#
# # 创建Redis的解析器对象
# redis_subparsers = subparsers.add_parser("redis", help="Redis相关漏扫")
# redis_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="Redis主机地址")
# redis_subparsers.add_argument("--port", "-P", dest="port", type=int, default=6379, help="Redis端口号")
# redis_subparsers.add_argument("--password", "-p", dest="password", type=argparse.FileType(mode='r', encoding="utf-8"),
#                               default="./passwd.dict", help="Redis密码本")
#
# # 创建IP的解析器对象
# ip_subparsers = subparsers.add_parser("ip", help="主机相关漏扫")
# ip_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="IP主机地址网段")
# ip_subparsers.add_argument("--iface", "-f", dest="iface", type=str, default="以太网", help="扫描网卡")
#
# # 创建Domain的解析器对象
# domain_subparsers = subparsers.add_parser("domain", help="子域名扫描")
# domain_subparsers.add_argument("--domain", "-d", dest="domain", type=str, required=True, help="要扫描的主域名")
# domain_subparsers.add_argument("--dict", "-D", dest="dict_file", type=argparse.FileType(mode='r', encoding="utf-8"),
#                               default="dic.txt", help="子域名字典文件")
#
# # 创建SSH的解析器对象
# ssh_subparsers = subparsers.add_parser("ssh", help="SSH相关漏扫")
# ssh_subparsers.add_argument("--host", "-H", dest="host", type=str, required=True, help="SSH主机地址")
# ssh_subparsers.add_argument("--port", "-P", dest="port", type=int, default=22, help="SSH端口号")
# ssh_subparsers.add_argument("--username", "-u", dest="username", type=argparse.FileType(mode='r', encoding="utf-8"),
#                             default="./user.dict", help="SSH用户名字典")
# ssh_subparsers.add_argument("--password", "-p", dest="password", type=argparse.FileType(mode='r', encoding="utf-8"),
#                             default="./passwd.dict", help="SSH密码字典")
#
# args = parser.parse_args()
#
# if args.sub_command == "mysql":
#     mysql_run(args.host, args.port, args.username.readlines(), args.password.readlines())
# elif args.sub_command == "redis":
#     redis_run(args.host, args.port, args.password.readlines())
# elif args.sub_command == "ip":
#     run_ip(args.host, args.iface)
# elif args.sub_command == "domain":
#     run_domain(args.domain, args.dict_file.name)
# elif args.sub_command == "ssh":  # 新增SSH命令处理
#     ssh_run(args.host, args.port, args.username.readlines(), args.password.readlines())
# else:
#     print("请选择要使用的功能")


import argparse
import threading
import warnings
import requests
from mysql_val import mysql_run
from redis_val import redis_run
from ip_val import run_ip
from ssh_val import ssh_run
from port_val import port_run  # 导入端口扫描功能
from ccolor import Colors

warnings.filterwarnings("ignore")

banner = """
        __                        .__
  _____|  | ___________  _______  |  |
 /  ___/  |/ / ____/\  \/ /\__  \ |  |
 \___ \|    < <_|  | \   /  / __ \|  |__
/____  >__|_ \__   |  \_/  (____  /____/
     \/     \/  |__|            \/
"""
print(f"{Colors.BLUE_BRIGHT}{Colors.BOLD}{banner}{Colors.RESET}")

headers = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/89.0.4389.90 Safari/537.36",
}

def subdomain_scan(domain, sub):
    host1 = "http://" + sub + "." + domain
    host2 = "https://" + sub + "." + domain
    try:
        response1 = requests.get(host1, headers=headers, timeout=5)
        if response1.status_code in [200,301,302,401,403,404]:
            print(f"[+] {sub}.{domain} 存活")
    except:
        try:
            response2 = requests.get(host2, headers=headers, timeout=5)
            if response2.status_code in [200,301,302,401,403,404]:
                print(f"[+] {sub}.{domain} 存活")
        except:
            pass

def run_domain(domain, sub_file="dic.txt"):
    threads = []
    with open(file=sub_file, mode="r", encoding="utf-8") as f:
        for sub in f:
            sub = sub.strip()
            t = threading.Thread(target=subdomain_scan, args=(domain, sub))
            t.start()
            threads.append(t)
    for t in threads:
        t.join()

# 创建命令行解析器对象
parser = argparse.ArgumentParser(description="智慧AI漏扫")
subparsers = parser.add_subparsers(help="选择要使用的功能", dest="sub_command")

# 创建MySQL的解析器对象
mysql_subparsers = subparsers.add_parser("mysql", help="MySQL相关漏扫")
mysql_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="MySQL主机地址")
mysql_subparsers.add_argument("--port", "-P", dest="port", type=int, default=3306, help="MySQL端口号")
mysql_subparsers.add_argument("--username", "-u", dest="username", type=argparse.FileType(mode='r', encoding="utf-8"),
                              default="./user.dict", help="MySQL用户名")
mysql_subparsers.add_argument("--password", "-p", dest="password", type=argparse.FileType(mode='r', encoding="utf-8"),
                              default="./passwd.dict", help="MySQL密码本")

# 创建Redis的解析器对象
redis_subparsers = subparsers.add_parser("redis", help="Redis相关漏扫")
redis_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="Redis主机地址")
redis_subparsers.add_argument("--port", "-P", dest="port", type=int, default=6379, help="Redis端口号")
redis_subparsers.add_argument("--password", "-p", dest="password", type=argparse.FileType(mode='r', encoding="utf-8"),
                              default="./passwd.dict", help="Redis密码本")

# 创建IP的解析器对象
ip_subparsers = subparsers.add_parser("ip", help="主机相关漏扫")
ip_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="IP主机地址网段")
ip_subparsers.add_argument("--iface", "-f", dest="iface", type=str, default="以太网", help="扫描网卡")

# 创建Domain的解析器对象
domain_subparsers = subparsers.add_parser("domain", help="子域名扫描")
domain_subparsers.add_argument("--domain", "-d", dest="domain", type=str, required=True, help="要扫描的主域名")
domain_subparsers.add_argument("--dict", "-D", dest="dict_file", type=argparse.FileType(mode='r', encoding="utf-8"),
                              default="./dic.txt", help="子域名字典文件")

# 创建SSH的解析器对象
ssh_subparsers = subparsers.add_parser("ssh", help="SSH相关漏扫")
ssh_subparsers.add_argument("--host", "-H", dest="host", type=str, required=True, help="SSH主机地址")
ssh_subparsers.add_argument("--port", "-P", dest="port", type=int, default=22, help="SSH端口号")
ssh_subparsers.add_argument("--username", "-u", dest="username", type=argparse.FileType(mode='r', encoding="utf-8"),
                            default="./user.dict", help="SSH用户名字典")
ssh_subparsers.add_argument("--password", "-p", dest="password", type=argparse.FileType(mode='r', encoding="utf-8"),
                            default="./passwd.dict", help="SSH密码字典")

# 创建Port的解析器对象
port_subparsers = subparsers.add_parser("port", help="端口扫描")
port_subparsers.add_argument("--host", "-H", dest="host", type=str, required=True, help="目标主机地址")
port_subparsers.add_argument("--scan-type", "-t", dest="scan_type", choices=["full", "common"],
                            default="common", help="扫描类型: full(全端口)或common(常见端口)")
port_subparsers.add_argument("--segment", "-s", dest="segment_size", type=int,
                            default=10, help="全端口扫描时的分段大小(仅对全扫描有效)")

args = parser.parse_args()

if args.sub_command == "mysql":
    mysql_run(args.host, args.port, args.username.readlines(), args.password.readlines())
elif args.sub_command == "redis":
    redis_run(args.host, args.port, args.password.readlines())
elif args.sub_command == "ip":
    run_ip(args.host, args.iface)
elif args.sub_command == "domain":
    run_domain(args.domain, args.dict_file.name)
elif args.sub_command == "ssh":
    ssh_run(args.host, args.port, args.username.readlines(), args.password.readlines())
elif args.sub_command == "port":
    port_run(args.host, args.scan_type, args.segment_size)
else:
    print("请选择要使用的功能")